from locale import normalize
from turtle import forward
import torch
import torch.nn as nn
import random
import scipy.io as scio
import numpy as np
from visdom import Visdom
import os
import matplotlib.pyplot as plt
import sys
import torch.nn.init as init
from torchsummary import summary
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import datetime
from tool import *
import torch.nn.functional as F
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)



#定义网络
class FC_Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(in_features=101,out_features=256)
        self.fc2 = nn.Linear(in_features=256,out_features=128)
        self.out = nn.Linear(in_features=128,out_features=101)

    def forward(self,t):
        t = t.reshape(-1,101)
        t = self.fc1(t)
        t = F.relu(t)
        t = self.fc2(t)
        t = F.relu(t)
        t = self.out(t)
        return t
#
class FC_IN_Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.IN = nn.LayerNorm(normalized_shape=101)
        self.fc1 = nn.Linear(in_features=101,out_features=256)
        self.fc2 = nn.Linear(in_features=256,out_features=128)
        self.out = nn.Linear(in_features=128,out_features=101)

    def forward(self,t):
        t = t.reshape(-1,101)
        #print(t)
        t = self.IN(t)
        #print(t)
        #print(t.shape)
        t = self.fc1(t)
        t = F.relu(t)
        t = self.fc2(t)
        t = F.relu(t)
        t = self.out(t)
        return t

class FC_IN_Net2(nn.Module):
    def __init__(self):
        super().__init__()
        self.IN = nn.LayerNorm(101)
        self.fc1 = nn.Linear(in_features=101,out_features=256)
        self.fc2 = nn.Linear(in_features=101,out_features=256)
        self.fc3 = nn.Linear(in_features=256,out_features=128)
        self.out = nn.Linear(in_features=128,out_features=101)

    def forward(self,t1,t2):
        t1 = t1.reshape(-1,101)
        t2 = t2.reshape(-1,101)
        t1 = self.IN(t1)
        t2 = self.IN(t2)
        t1 = self.fc1(t1)
        t2 = self.fc2(t2)
        t = F.relu(t1+t2)
        t = self.fc3(t)
        t = F.relu(t)
        t = self.out(t)
        #t = t.reshape(4,101,-1)
        return t


class FC_Net2(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(in_features=101,out_features=256)
        self.fc2 = nn.Linear(in_features=101,out_features=256)
        self.fc3 = nn.Linear(in_features=256,out_features=128)
        self.out = nn.Linear(in_features=128,out_features=101)

    def forward(self,t1,t2):
        t1 = t1.reshape(-1,101)
        t2 = t2.reshape(-1,101)
        t1 = self.fc1(t1)
        t2 = self.fc2(t2)
        t = F.relu(t1+t2)
        t = self.fc3(t)
        t = F.relu(t)
        t = self.out(t)
        #t = t.reshape(4,101,-1)
        return t

class CNN_Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv2d = nn.Conv1d(in_channels=1,out_channels=8,kernel_size=1,stride=1,padding=0)
        self.conv2d1 = nn.Conv1d(in_channels=8,out_channels=1,kernel_size=1,stride=1,padding=0)
        self.fc1 = nn.Linear(in_features=101,out_features=101)

    def forward(self,t):
        t = self.conv2d(t)
        
        t = F.relu(t)
        t = self.conv2d1(t)
        
        t = F.relu(t)
        t = self.fc1(t)
        return t


# 初始化权重
def weights_init(m):
    if type(m) in [nn.ConvTranspose2d, nn.Conv2d]:
        init.xavier_normal_(m.weight)
    elif type(m) == nn.BatchNorm2d:
        init.normal_(m.weight, 1.0, 0.02)
        init.constant_(m.bias, 0)



transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片转换为Tensor,归一化至[0,1]
])

class MuscleData1(Dataset):
    #第一类数据，输入100x1的data和100x1的label
    def __init__(self, data, label):
        self.data = data
        self.label = label
        self.transforms = transform
        
        
    def __getitem__(self,index):
        muscle_data = self.data[index]
        muscle_label = self.label[index]
        #muscle_data = np.squeeze(muscle_data)
        #muscle_label = np.squeeze(muscle_label)
        muscle_data =muscle_data.astype(np.float32)
        muscle_label =muscle_label.astype(np.float32)
        #muscle_data = self.transforms(muscle_data)
        #print(muscle_label,muscle_data)
        return muscle_data,muscle_label

    def __len__(self): #len不可缺省。Dataset返回的数据若没有len则变成列表。即调用时候需要set[index]。
        return len(self.label)


class MuscleData2(Dataset):
    #第一类数据，输入100x1的data和100x1的label
    def __init__(self, data1,data2,label):
        self.data1 = data1
        self.data2 =data2
        self.label = label
        self.transforms = transform
        
        
    def __getitem__(self,index):
        muscle_data1 = self.data1[index]
        muscle_data2 = self.data2[index]
        muscle_label = self.label[index]
        #muscle_data = np.squeeze(muscle_data)
        #muscle_label = np.squeeze(muscle_label)
        muscle_data1 =muscle_data1.astype(np.float16)
        muscle_data2 =muscle_data2.astype(np.float16)
        muscle_label =muscle_label.astype(np.float16)
        #muscle_data = self.transforms(muscle_data)
        #print(muscle_label,muscle_data)
        return muscle_data1,muscle_data2,muscle_label

    def __len__(self): #len不可缺省。Dataset返回的数据若没有len则变成列表。即调用时候需要set[index]。
        return len(self.label)